#include "UnstructMG.hpp"
#include <chrono>
#include <random>
#include <algorithm>
#include <vector>
#include <string>

#include "type_header.h"
using idx_t = int;

class TEST_RECORD {
public:
    double setup = 0.0, solve = 0.0;
    double prec = 0.0, oper = 0.0, axpy = 0.0, dot = 0.0;
    int iter = 0;
    bool operator < (const TEST_RECORD & rec) const {
        double my_tot = setup + solve;
        double his_tot = rec.setup + rec.solve;
        return my_tot < his_tot;
    }
};

struct TEST_CONFIG
{
    int max_iter = 100;
    double atol = 0.0;// 绝对残差
    double rtol = 1.0e-9;// 相对残差
    int iter_print_level = 2;
    int prec_print_level = 1;
    int restart_len = 10;
    std::string config_mg_file;
};

template<typename idx_t>
class PartDescriptor
{
public:
    idx_t id;// 全局块序号
    idx_t ilower[3], iupper[3];
    idx_t nelems;
};

#ifndef DOF
#define DOF 1
#endif

IterativeSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, DOF> * solver  = nullptr;
Solver         <idx_t,     ksp_t, pc_calc_t       , DOF> * precond = nullptr;

double  min_times[NUM_KRYLOV_RECORD], 
        max_times[NUM_KRYLOV_RECORD],
        avg_times[NUM_KRYLOV_RECORD];

void buildup_solver(std::string solver_name, std::string prec_name, const TEST_CONFIG & config_file);
void destroy_solver();
void setup_and_solve(const par_CSRMatrix<idx_t, ksp_t, ksp_t, DOF> & ss_A,
    const par_Vector<idx_t, ksp_t, DOF> & b, par_Vector<idx_t, ksp_t, DOF> & x,
    TEST_RECORD & record);

void check_residual(const par_CSRMatrix<idx_t, ksp_t, ksp_t, DOF> & ss_A,
    const par_Vector<idx_t, ksp_t, DOF> & ss_x, const par_Vector<idx_t, ksp_t, DOF> & ss_b,
    par_Vector<idx_t, ksp_t, DOF> & ss_y, double & r_nrm2, double & b_nrm2);
void stat_part_times(TEST_RECORD & record);
void stat_multi_runs(std::vector<TEST_RECORD> & records);

void buildup_solver(std::string its_name, std::string prc_name, const TEST_CONFIG & config_file)
{
    assert(solver == nullptr);
    assert(precond== nullptr);
    
    int my_pid; MPI_Comm_rank(MPI_COMM_WORLD, &my_pid);
    if (my_pid == 0) {
        printf("Bytes : ksp_t %ld pc_calc_t %ld pc_data_t %ld\n", sizeof(ksp_t), sizeof(pc_calc_t), sizeof(pc_data_t));
    }
    const int restart = config_file.restart_len;
    // 设置求解器
    if (its_name == "GCR") {
        using GCR = GCRSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, DOF>;
        solver = new GCR();
        ((GCR*)solver)->SetInnerIterMax(restart);
    } else if (its_name == "CG") {
        solver = new CGSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, DOF>;
    } else if (its_name == "GMRES") {
        using GMRES = GMRESSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, DOF>;
        solver = new GMRES();
        ((GMRES*)solver)->SetRestartlen(restart);
    } else if (its_name == "BiCG") {
        using BiCG = BiCGStabSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, DOF>;
        solver = new BiCG();
    } else {
        if (my_pid == 0) printf("INVALID iterative solver name of %s\nOnly GCR, CG, GMRES, FGMRES available\n", its_name.c_str());
        MPI_Abort(MPI_COMM_WORLD, -1);
    }
    solver->SetMaxIter(config_file.max_iter);
    if (its_name == "GCR") {// 里面用的范数没有开方
        solver->SetRelTol(config_file.rtol * config_file.rtol);
        solver->SetAbsTol(config_file.atol * config_file.atol);
    } else {
        solver->SetRelTol(config_file.rtol);
        solver->SetAbsTol(config_file.atol);
    }
    
    // 设置预条件子
    if (strstr(prc_name.c_str(), "PJ")) {
        if (my_pid == 0) printf("  using \033[1;35mpointwise-Jacobi\033[0m as preconditioner\n");
        precond = new PointJacobi<idx_t, pc_data_t, ksp_t, pc_calc_t, DOF>();
    } else if (strstr(prc_name.c_str(), "PGS")) {
        using PGS = PointGS<idx_t, pc_data_t, ksp_t, pc_calc_t, DOF>;
        PGS::SCAN_TYPE type = PGS::SYMMETRIC;				
        if      (strstr(prc_name.c_str(), "F")) type = PGS::FORWARD;
        else if (strstr(prc_name.c_str(), "B")) type = PGS::BACKWARD;
        else if (strstr(prc_name.c_str(), "S")) type = PGS::SYMMETRIC;
        if (my_pid == 0) printf("  using \033[1;35mpointwise-GS %d\033[0m as preconditioner\n", type);
        precond = new PGS(type);
    } else if (prc_name == "ILU") {
        using ILU = BoxILU<idx_t, pc_data_t, ksp_t, pc_calc_t, DOF>;
        precond = new ILU;
    } else if (prc_name == "ILUT") {
        using ILU = BoxILU<idx_t, pc_data_t, ksp_t, pc_calc_t, DOF>;
        precond = new ILU(1);
    } else if (prc_name == "GMG") {
        using MG = UnstructMG<idx_t, pc_data_t, ksp_t, pc_calc_t, DOF>;
        precond = new MG(config_file.config_mg_file);
    }
    else 
    {
        if (my_pid == 0) printf("No Precond!\n");
    }

    if (precond != nullptr) solver->SetPreconditioner(*precond);
}

void setup_and_solve(const par_CSRMatrix<idx_t, ksp_t, ksp_t, DOF> & ss_A,
    const par_Vector<idx_t, ksp_t, DOF> & b, par_Vector<idx_t, ksp_t, DOF> & x,
    TEST_RECORD & record)
{
    MPI_Barrier(MPI_COMM_WORLD);
    double t0 = wall_time();
    solver->SetOperator(ss_A);
    t0 = wall_time() - t0;

    MPI_Barrier(MPI_COMM_WORLD);
    double t1 = wall_time();
    solver->Mult(b, x, false);
    t1 = wall_time() - t1;

    record.setup = (precond != nullptr) ? precond->setup_time : 0.0;
    record.solve = t1;
    record.iter  = solver->GetNumIterations();
}

double check_eigen_value(idx_t max_iter, const par_CSRMatrix<idx_t, ksp_t, ksp_t, DOF> & A)
{
    return A.power_iteration(max_iter, nullptr);
}

double check_preconditioned_eigen_value(idx_t max_iter, const par_CSRMatrix<idx_t, ksp_t, ksp_t, DOF> & A)
{
    return A.power_iteration(max_iter, (const Operator<idx_t, ksp_t, DOF>*) precond);
}

void destroy_solver()
{
    if (precond) { delete precond; precond = nullptr; }
    if (solver ) { delete solver ; solver  = nullptr; }
}

void check_residual(const par_CSRMatrix<idx_t, ksp_t, ksp_t, DOF> & ss_A,
    const par_Vector<idx_t, ksp_t, DOF> & ss_x, const par_Vector<idx_t, ksp_t, DOF> & ss_b,
    par_Vector<idx_t, ksp_t, DOF> & ss_y, double & r_nrm2, double & b_nrm2)
{
    b_nrm2 = vec_dot<idx_t, ksp_t, double>(ss_b, ss_b);
    b_nrm2 = sqrt(b_nrm2);

    ss_A.Mult(ss_x, ss_y, false);
    vec_add(ss_b, (ksp_t) -1.0, ss_y, ss_y);
    r_nrm2 = vec_dot<idx_t, ksp_t, double>(ss_y, ss_y);
    r_nrm2 = sqrt(r_nrm2);
}

void stat_part_times(TEST_RECORD & record)
{
    int my_pid; MPI_Comm_rank(MPI_COMM_WORLD, &my_pid);
    int num_procs; MPI_Comm_size(MPI_COMM_WORLD, &num_procs);
    MPI_Allreduce(solver->part_times, min_times, NUM_KRYLOV_RECORD, MPI_DOUBLE, MPI_MIN, MPI_COMM_WORLD);
    MPI_Allreduce(solver->part_times, max_times, NUM_KRYLOV_RECORD, MPI_DOUBLE, MPI_MAX, MPI_COMM_WORLD);
    MPI_Allreduce(solver->part_times, avg_times, NUM_KRYLOV_RECORD, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD); 
    for (int i = 0; i < NUM_KRYLOV_RECORD; i++)
        avg_times[i] /= num_procs;
    if (my_pid == 0) {
        printf("prec time min/avg/max %.3e %.3e %.3e\n", min_times[PREC], avg_times[PREC], max_times[PREC]);
        printf("oper time min/avg/max %.3e %.3e %.3e\n", min_times[OPER], avg_times[OPER], max_times[OPER]);
        printf("axpy time min/avg/max %.3e %.3e %.3e\n", min_times[AXPY], avg_times[AXPY], max_times[AXPY]);
        printf("dot  tune min/avg/max %.3e %.3e %.3e\n", min_times[DOT ], avg_times[DOT ], max_times[DOT ]);
    }

    record.prec = avg_times[PREC];
    record.oper = avg_times[OPER];
    record.axpy = avg_times[AXPY];
    record.dot  = avg_times[DOT ];
}

void stat_multi_runs(std::vector<TEST_RECORD> & records)
{
    const idx_t test_cnt = records.size();
    std::sort(records.begin(), records.end());
    const TEST_RECORD & best = records[0], & worst = records[test_cnt - 1];
    TEST_RECORD avg;
    for (idx_t i = 0; i < test_cnt; i++) {
        avg.setup += records[i].setup;
        avg.solve += records[i].solve;
        avg.prec  += records[i].prec;
        avg.iter  += records[i].iter;
    }
    avg.setup /= test_cnt;
    avg.solve /= test_cnt;
    avg.prec  /= test_cnt;
    avg.iter  /= test_cnt;
    double min_tot = best .setup + best .solve;
    double max_tot = worst.setup + worst.solve;
    double avg_tot = avg  .setup + avg  .solve;
    printf("\n  Summary of %d tests:\n", test_cnt);
    printf("     Setup time  Solve time  #Iter  Total time\n");
    printf("Min    %.6f    %.6f  %5d  %.6f\n", best .setup, best .solve, best .iter, min_tot);
    printf("Avg    %.6f    %.6f  %5d  %.6f\n", avg  .setup, avg  .solve, avg  .iter, avg_tot);
    printf("Max    %.6f    %.6f  %5d  %.6f\n", worst.setup, worst.solve, worst.iter, max_tot);    
}